package org.yinweichen.llm_dataset_backend.util;

import java.util.List;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Getter;
import lombok.Setter;

import java.io.File;
import java.io.IOException;

public class QAReader {
    @Setter
    @Getter
    public static class OriginalQA {
        // Getters and setters
        @JsonProperty("source_question_id")
        private int sourceQuestionId;

        @JsonProperty("source_answer_id")
        private int sourceAnswerId;

        @JsonProperty("question_title")
        private String questionTitle;

        @JsonProperty("question_body")
        private String questionBody;

        @JsonProperty("answer_body")
        private String answerBody;

        @JsonProperty("source")
        private String source;

        @JsonProperty("source_id")
        private int sourceId;

        @JsonProperty("source_url")
        private String sourceUrl;

    }

    @Setter
    @Getter
    public static class StandardQA {
        // Getters and setters
        @JsonProperty("qa_id")
        private int qaId;

        @JsonProperty("source_question_id")
        private int sourceQuestionId;

        @JsonProperty("source_answer_id")
        private int sourceAnswerId;

        @JsonProperty("question")
        private String question;

        @JsonProperty("answer")
        private String answer;

        @JsonProperty("key_points")
        private List<String> keyPoints;

    }

    public static List<OriginalQA> loadOriginalQA(String filePath) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        return mapper.readValue(
                new File(filePath),
                mapper.getTypeFactory().constructCollectionType(List.class, OriginalQA.class)
        );
    }

    public static List<StandardQA> loadStandardQA(String filePath) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        return mapper.readValue(
                new File(filePath),
                mapper.getTypeFactory().constructCollectionType(List.class, StandardQA.class)
        );
    }
}
